#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
创建高质量的语音助手对话短文本语义匹配数据集

Author: BOSS (牛马)
Date: 2024-06-19
"""

import pandas as pd
import numpy as np
import random
from pathlib import Path
import itertools

class VoiceAssistantDatasetGenerator:
    """语音助手数据集生成器"""
    
    def __init__(self):
        # 语音助手常见功能类别
        self.categories = {
            "音乐控制": {
                "base_intents": [
                    "播放音乐", "暂停音乐", "停止播放", "下一首", "上一首", 
                    "调节音量", "静音", "播放歌单", "搜索歌曲", "收藏音乐"
                ],
                "variations": [
                    "放音乐", "听音乐", "来首歌", "播歌", "放歌",
                    "暂停", "停止", "停下", "别播了", "关掉音乐",
                    "声音大点", "声音小点", "调大音量", "调小音量", "音量调节",
                    "下一个", "切歌", "换一首", "下首歌", "跳过",
                    "上一个", "前一首", "回到上一首", "返回上一曲"
                ]
            },
            
            "通讯功能": {
                "base_intents": [
                    "打电话", "拨打电话", "呼叫", "发短信", "发消息", 
                    "发微信", "查看通讯录", "添加联系人", "视频通话", "语音通话"
                ],
                "variations": [
                    "给某某打电话", "拨号", "联系", "通话", "呼叫某某",
                    "发信息", "发条短信", "发个消息", "短信", "消息",
                    "发微信消息", "微信", "发朋友圈", "微信聊天",
                    "视频", "视频聊天", "开视频", "语音聊天"
                ]
            },
            
            "天气查询": {
                "base_intents": [
                    "查看天气", "天气预报", "今天天气", "明天天气", "天气情况",
                    "温度", "下雨吗", "需要带伞吗", "穿什么衣服", "空气质量"
                ],
                "variations": [
                    "天气怎么样", "今天天气如何", "外面天气", "天气状况",
                    "今天几度", "温度多少", "热不热", "冷不冷",
                    "会下雨吗", "有雨吗", "要下雨了吗", "雨天",
                    "空气好吗", "PM2.5", "雾霾", "空气指数"
                ]
            },
            
            "导航出行": {
                "base_intents": [
                    "导航", "路线规划", "怎么走", "去某地", "回家",
                    "附近", "找地方", "定位", "打车", "叫车"
                ],
                "variations": [
                    "导航到", "去哪里", "路线", "怎么去", "如何到达",
                    "回家路线", "回公司", "回学校", "回家怎么走",
                    "附近的", "周边", "找个", "哪里有", "最近的",
                    "滴滴", "出租车", "网约车", "叫个车"
                ]
            },
            
            "时间日程": {
                "base_intents": [
                    "几点了", "现在时间", "设置闹钟", "定闹钟", "提醒",
                    "日程安排", "查看日历", "添加日程", "会议提醒", "倒计时"
                ],
                "variations": [
                    "现在几点", "时间", "几点", "什么时候",
                    "闹钟", "叫醒我", "明天叫我", "定时",
                    "提醒我", "别忘了", "记住", "备忘",
                    "日历", "安排", "计划", "行程", "日程表"
                ]
            },
            
            "设备控制": {
                "base_intents": [
                    "开灯", "关灯", "调节亮度", "开空调", "关空调",
                    "调温度", "开电视", "关电视", "换台", "调节音量"
                ],
                "variations": [
                    "打开灯", "关闭灯光", "灯光", "照明", "亮一点",
                    "空调", "制冷", "制热", "温度", "热一点", "冷一点",
                    "电视", "看电视", "关电视", "电视机", "换个台"
                ]
            },
            
            "信息查询": {
                "base_intents": [
                    "搜索", "查询", "百度一下", "查找", "搜一下",
                    "新闻", "股票", "汇率", "翻译", "计算"
                ],
                "variations": [
                    "搜索一下", "查一查", "找一找", "百度", "谷歌",
                    "看新闻", "新闻资讯", "今日新闻", "头条",
                    "股价", "股市", "基金", "理财", "投资",
                    "翻译成", "英文", "中文", "日语", "韩语",
                    "算一下", "计算器", "多少", "等于"
                ]
            },
            
            "娱乐休闲": {
                "base_intents": [
                    "讲笑话", "讲故事", "唱歌", "聊天", "玩游戏",
                    "看视频", "看电影", "听广播", "听书", "冥想"
                ],
                "variations": [
                    "说个笑话", "来个笑话", "逗我笑", "幽默一下",
                    "故事", "童话", "小说", "讲个故事",
                    "唱首歌", "来首歌", "音乐", "歌曲",
                    "陪我聊天", "说话", "对话", "交流",
                    "游戏", "小游戏", "娱乐", "玩一下"
                ]
            }
        }
        
        # 完全不相关的对话对
        self.unrelated_pairs = [
            ("播放音乐", "查看天气"),
            ("打电话", "开灯"),
            ("设置闹钟", "搜索餐厅"),
            ("导航回家", "讲笑话"),
            ("发短信", "调节温度"),
            ("查看新闻", "播放视频"),
            ("翻译", "打车"),
            ("计算", "唱歌"),
            ("股票查询", "关电视"),
            ("空气质量", "玩游戏")
        ]
    
    def generate_similar_pairs(self, n_pairs=500):
        """生成相似语义对"""
        similar_pairs = []
        
        for category, data in self.categories.items():
            base_intents = data["base_intents"]
            variations = data["variations"]
            
            # 基础意图与变体的配对
            for base_intent in base_intents:
                for variation in variations:
                    if self._is_semantically_related(base_intent, variation):
                        similar_pairs.append((base_intent, variation, 1))
            
            # 同类意图之间的相似配对
            for i, intent1 in enumerate(base_intents):
                for intent2 in base_intents[i+1:]:
                    if self._is_semantically_similar(intent1, intent2):
                        similar_pairs.append((intent1, intent2, 1))
            
            # 变体之间的配对
            for i, var1 in enumerate(variations):
                for var2 in variations[i+1:]:
                    if self._is_semantically_similar(var1, var2):
                        similar_pairs.append((var1, var2, 1))
        
        # 随机采样
        if len(similar_pairs) > n_pairs:
            similar_pairs = random.sample(similar_pairs, n_pairs)
        
        return similar_pairs
    
    def generate_dissimilar_pairs(self, n_pairs=500):
        """生成不相似语义对"""
        dissimilar_pairs = []
        
        # 跨类别的不相似对
        categories = list(self.categories.keys())
        for i, cat1 in enumerate(categories):
            for cat2 in categories[i+1:]:
                intents1 = self.categories[cat1]["base_intents"] + self.categories[cat1]["variations"]
                intents2 = self.categories[cat2]["base_intents"] + self.categories[cat2]["variations"]
                
                # 随机组合
                for _ in range(min(50, len(intents1) * len(intents2) // 10)):
                    intent1 = random.choice(intents1)
                    intent2 = random.choice(intents2)
                    if not self._is_semantically_related(intent1, intent2):
                        dissimilar_pairs.append((intent1, intent2, 0))
        
        # 添加预定义的不相关对
        for pair in self.unrelated_pairs:
            dissimilar_pairs.append((pair[0], pair[1], 0))
            dissimilar_pairs.append((pair[1], pair[0], 0))  # 反向也加入
        
        # 随机采样
        if len(dissimilar_pairs) > n_pairs:
            dissimilar_pairs = random.sample(dissimilar_pairs, n_pairs)
        
        return dissimilar_pairs
    
    def _is_semantically_related(self, text1, text2):
        """判断两个文本是否语义相关"""
        # 简单的关键词匹配
        keywords1 = set(text1.replace("某某", "").split())
        keywords2 = set(text2.replace("某某", "").split())
        
        # 如果有共同关键词，认为相关
        common_keywords = keywords1 & keywords2
        if common_keywords:
            return True
        
        # 特定的语义关联规则
        semantic_groups = [
            {"音乐", "歌", "播放", "听", "放"},
            {"电话", "通话", "拨打", "呼叫", "联系"},
            {"天气", "温度", "雨", "热", "冷"},
            {"导航", "路线", "去", "走", "到"},
            {"时间", "点", "闹钟", "提醒"},
            {"灯", "亮", "照明", "开", "关"},
            {"搜索", "查", "找", "百度"},
        ]
        
        for group in semantic_groups:
            if any(keyword in text1 for keyword in group) and any(keyword in text2 for keyword in group):
                return True
        
        return False
    
    def _is_semantically_similar(self, text1, text2):
        """判断两个文本是否语义相似"""
        if text1 == text2:
            return False
        
        return self._is_semantically_related(text1, text2)
    
    def add_noise_and_variations(self, pairs):
        """添加噪声和变体"""
        enhanced_pairs = []
        
        for query1, query2, label in pairs:
            # 原始对
            enhanced_pairs.append((query1, query2, label))
            
            # 添加语气词和修饰词
            prefixes = ["请", "帮我", "我想", "能否", "可以", "麻烦"]
            suffixes = ["吧", "一下", "呢", "啊", "哦", ""]
            
            # 随机添加前缀和后缀
            if random.random() < 0.3:
                query1_enhanced = random.choice(prefixes) + query1 + random.choice(suffixes)
                enhanced_pairs.append((query1_enhanced, query2, label))
            
            if random.random() < 0.3:
                query2_enhanced = random.choice(prefixes) + query2 + random.choice(suffixes)
                enhanced_pairs.append((query1, query2_enhanced, label))
        
        return enhanced_pairs
    
    def generate_dataset(self, train_size=2000, test_size=500, positive_ratio=0.5):
        """生成完整数据集"""
        print("🔄 生成语音助手语义匹配数据集...")
        
        # 计算正负样本数量
        train_positive = int(train_size * positive_ratio)
        train_negative = train_size - train_positive
        test_positive = int(test_size * positive_ratio)
        test_negative = test_size - test_positive
        
        # 生成训练集
        print(f"   生成训练集: {train_size} 样本 (正样本: {train_positive}, 负样本: {train_negative})")
        train_similar = self.generate_similar_pairs(train_positive)
        train_dissimilar = self.generate_dissimilar_pairs(train_negative)
        
        # 生成测试集
        print(f"   生成测试集: {test_size} 样本 (正样本: {test_positive}, 负样本: {test_negative})")
        test_similar = self.generate_similar_pairs(test_positive)
        test_dissimilar = self.generate_dissimilar_pairs(test_negative)
        
        # 合并并打乱
        train_data = train_similar + train_dissimilar
        test_data = test_similar + test_dissimilar
        
        # 添加变体和噪声
        train_data = self.add_noise_and_variations(train_data)
        test_data = self.add_noise_and_variations(test_data)
        
        # 重新采样到目标大小
        train_data = random.sample(train_data, min(len(train_data), train_size))
        test_data = random.sample(test_data, min(len(test_data), test_size))
        
        # 打乱数据
        random.shuffle(train_data)
        random.shuffle(test_data)
        
        # 创建DataFrame
        train_df = pd.DataFrame(train_data, columns=['query1', 'query2', 'label'])
        train_df['id'] = range(len(train_df))
        train_df = train_df[['id', 'query1', 'query2', 'label']]
        
        # 测试集不包含标签
        test_df = pd.DataFrame(test_data, columns=['query1', 'query2', 'label'])
        test_df = test_df[['query1', 'query2']].copy()
        test_df['id'] = range(len(test_df))
        test_df = test_df[['id', 'query1', 'query2']]
        
        return train_df, test_df
    
    def save_dataset(self, train_df, test_df, output_dir="data/raw"):
        """保存数据集"""
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        
        train_path = output_dir / "train.csv"
        test_path = output_dir / "test.csv"
        
        train_df.to_csv(train_path, index=False)
        test_df.to_csv(test_path, index=False)
        
        print(f"✅ 数据集已保存:")
        print(f"   训练集: {train_path} ({len(train_df)} 样本)")
        print(f"   测试集: {test_path} ({len(test_df)} 样本)")
        
        # 显示标签分布
        label_dist = train_df['label'].value_counts().sort_index()
        print(f"   训练集标签分布: {dict(label_dist)}")
        
        return train_path, test_path

def main():
    """主函数"""
    print("🤖 语音助手数据集生成器")
    print("Author: BOSS (牛马)")
    print("=" * 60)
    
    # 创建生成器
    generator = VoiceAssistantDatasetGenerator()
    
    # 生成数据集
    train_df, test_df = generator.generate_dataset(
        train_size=2000,
        test_size=500,
        positive_ratio=0.5
    )
    
    # 保存数据集
    train_path, test_path = generator.save_dataset(train_df, test_df)
    
    # 显示样本
    print(f"\n📊 训练集样本预览:")
    print(train_df.head(10))
    
    print(f"\n📊 测试集样本预览:")
    print(test_df.head(5))
    
    print(f"\n🎉 数据集生成完成!")
    print(f"📝 下一步: 运行 python demo_standalone.py 开始训练")

if __name__ == "__main__":
    main()
